#!/usr/bin/env python3
'''
未验证

训练记录：
在笔记本上训练
20241124：训练分数167分，测试分数327分，继续训练
20241124：训练分数210分，测试分数520分，继续训练
20241125：训练分数221分，测试分数925分，继续训练
20241126：训练分数228分，测试分数1624分，继续训练
20241127：训练分数256分，测试分数1624（有产生新的测试分数），继续训练，以及是否需要调整学习率
20241128: 训练分数299分，测试分数3281分，继续训练
20241129: 训练分数位置，测试分数3281（有产生新的最好分数2600+），继续训练
20241130: 训练分数336分，测试分数3281（有产生新的最好分数2700+），继续训练
20241202：训练分数327分，测试分数3281（产生新的最好分数2031+），调整学习率，继续训练
20241203: 训练分数325分，测试分数3281（新分数2869分），评估学习率，继续训练
20241204：学习率看起来没什么问题，此时为8e-5，变化两次，尝试增大训练缓冲区的大小
20241204：训练分数329分，测试分数3281（新分数2118分），貌似没有进步，评估是否需要继续训练
20241205: 学习率是6e-5，看起来没什么问题，变化两次，但是训练分数和测试分数都没有进步？到dqn basic的极限了吗？继续训练一天，无任何进步则不再继续训练，训练分数：425分，测试分数3281（新分数2690），继续训练
20241206:学习率是5e-5，学习率无问题，训练分数：346分，测试分数：3281（新分数：2368分），评估学习率是否暂停训练
20241207：学习率是4.3-05，训练分数：396分，测试分数：5829分，评估学习率继续训练
20241208: 学习率是 3.486-05，训练分数：373分（初始分数有732分，感觉不作数），测试分数：5829分，评估学习率是否停止训练？感觉到极限了吧
20241209: 继续训练一天，无进步则停止；学习率2.82e-05,训练分数：359，测试分数6243分，继续训练
20241210: 学习率2.5418658283290016e-05，有效训练分数：363分，测试分数分数：6243分，继续训练
20241211: 学习率2.2876792454961016e-05，有效训练分数：361分，测试分数6243分，无进步，停止训练，可能已经到头，也可能学习率等超参数存在问题
'''
import os
import gymnasium as gym
import ptan
import argparse

import logging
from logging.handlers import RotatingFileHandler
import torch
import torch.optim as optim
import torch.nn.functional as F
import numpy as np

from tensorboardX import SummaryWriter
from typing import Any
from lib import dqn_model, common
import ale_py

gym.register_envs(ale_py)


class FireResetEnv(gym.Wrapper):
    def __init__(self, env=None):
        """For environments where the user need to press FIRE for the game to start."""
        super(FireResetEnv, self).__init__(env)
        # 以下可知，一些游戏存在FIRE的动作，并且存在FIRE动作的游戏其游戏动作执行有三个以上
        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'
        assert len(env.unwrapped.get_action_meanings()) >= 3

    def step(self, action):
        return self.env.step(action)

    def reset(self, seed: int | None = None, options: dict[str, Any] | None = None):
        # 这里之所以尝试重置后尝试各种动作，是因为不知道哪个是FIRE，继续游戏，所以一个一个尝试
        # 如果不小心游戏结束了，则继续重置
        # 假设游戏继续游戏的按钮在前3
        self.env.reset(seed=seed, options=options)
        obs, _, done, _, info = self.env.step(1)
        if done:
            self.env.reset(seed=seed, options=options)
        obs, _, done, _, info = self.env.step(2)
        if done:
            self.env.reset(seed=seed, options=options)
        return obs, info

class RewardPenaltyWrapper(gym.Wrapper):
    def __init__(self, env, frame_penalty=-0.1, life_loss_penalty=-10):
        super(RewardPenaltyWrapper, self).__init__(env)
        self.frame_penalty = frame_penalty
        self.life_loss_penalty = life_loss_penalty
        self.previous_lives = 0

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        self.previous_lives = info.get('lives', 0)  # 初始生命值
        return obs, info

    def step(self, action):
        obs, reward, done, truncated, info = self.env.step(action)

        reward /= 100 # 缩放奖励
        
        # 处理生命减少时的惩罚
        current_lives = info.get('lives', self.previous_lives)
        if current_lives < self.previous_lives:
            reward += self.life_loss_penalty
            self.previous_lives = current_lives
        
        return obs, reward, done, truncated, info
    


def wrap_dqn(env, stack_frames=4, episodic_life=True, reward_clipping=True):
    if episodic_life:
        # 将多条生命的游戏模拟成单条生命ActorCriticAgent
        env = ptan.common.wrappers.EpisodicLifeEnv(env)
    # 增强初始化
    env = ptan.common.wrappers.NoopResetEnv(env, noop_max=30)

    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = FireResetEnv(env)
    env = ptan.common.wrappers.ProcessFrame84(env)
    env = ptan.common.wrappers.ImageToPyTorch(env)
    env = ptan.common.wrappers.FrameStack(env, stack_frames)
    env = RewardPenaltyWrapper(env)
    return env


def test_model(env, net, device, episodes=5):
    total_reward = 0.0
    for _ in range(episodes):
        noop_action_count = 0
        pre_action = -1
        obs, _ = env.reset()
        while True:
            obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
            logits_v = net(obs_v)
            probs_v = F.softmax(logits_v, dim=1)
            probs = probs_v.data.cpu().numpy()
            action = np.argmax(probs)
            if action == 0 and pre_action == action:  # Noop
                noop_action_count += 1
                if noop_action_count > 30:
                    break
            else:
                noop_action_count = 0
            pre_action = action
            obs, reward, done, trunc, _ = env.step(action)
            total_reward += reward
            if done or trunc:
                break
    return total_reward / episodes


def setup_logger(save_path):
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    handler = RotatingFileHandler(os.path.join(save_path, 'train.log'), maxBytes=1024 * 1024, backupCount=2)
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    return logger


if __name__ == "__main__":
    params = {
        'env_name':         "ALE/Atlantis-v5",
        'stop_reward':      5000.0,
        'run_name':         'dqn-basic-atlantis',
        'replay_size':      10 ** 6,
        'replay_initial':   50000,
        'target_net_sync':  10000,
        'epsilon_frames':   10 ** 6,
        'epsilon_start':    1.0,
        'epsilon_final':    0.1,
        'learning_rate':    0.00025,
        'gamma':            0.99,
        'batch_size':       64
    }

    save_path = os.path.join("saves", params['run_name'])
    os.makedirs(save_path, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    env = wrap_dqn(gym.make(params['env_name'], obs_type='rgb', frameskip=4, repeat_action_probability=0.0), episodic_life=False)
    test_env = wrap_dqn(gym.make(params['env_name'], obs_type='rgb', frameskip=4, repeat_action_probability=0.0), episodic_life=False)


    # 创建tensor board的存储目录
    writer = SummaryWriter(comment="atlantis-basic")
    # 创建dqn网络模型
    net = dqn_model.DQNBreakOut(env.observation_space.shape, env.action_space.n).to(device)

    # 考虑dqn网络，创建一个目标网络,负责计算在线网络net的目标q值
    tgt_net = ptan.agent.TargetNet(net)
    # 创建一个基于epsilon进行动作动作选择的选择器
    selector = ptan.actions.EpsilonGreedyActionSelector(epsilon=1.0)
    # epsion的值含义:epsion标识一个会随着训练的增加而逐渐减少的随机值
    # 表示训练一开始以随机选择执行的动作为主，因为此时网络并没有训练，进行随机探索
    # 随着训练的增加，网络已经具备一些情况下的决策能力，可以进行动作的选择
    epsilon_tracker = common.EpsilonTracker(selector, 1.0, 0.1, 10**6)
    # 创建一个dqn推理网络的代理
    agent = ptan.agent.DQNAgent(net, selector, device=device)

    # 创建经验重放缓冲区（也就是训练过程中采集的样本）
    # ExperienceSourceFirstLast应该是存储样本的缓存去
    exp_source = ptan.experience.ExperienceSourceFirstLast(env, agent, gamma=0.99, steps_count=1)
    # 第二个应该是从缓存区提取一小段样本进行训练的
    buffer = ptan.experience.ExperienceReplayBuffer(exp_source, buffer_size=3*10 ** 5)

    # 优化器
    optimizer = optim.Adam(net.parameters(), lr=0.0001)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=600000, gamma=0.9)

    # 统计所执行经理的帧数（轮数）
    frame_idx = 0
    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)),
                             key=lambda x: int(x.split('_')[2].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        net.load_state_dict(checkpoint['net'])
        tgt_net.target_model.load_state_dict(checkpoint['tgt_net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        frame_idx = checkpoint['frame_idx']
        if "scheduler" in checkpoint:
            scheduler.load_state_dict(checkpoint['scheduler'])
        print("加载模型成功")
        print("学习率：", optimizer.param_groups[0]['lr'])
        print("frame_idx: ", frame_idx)
        print("scheduler last epoch: ", scheduler.last_epoch)

    logger = setup_logger(save_path)
    # 创建一个奖励跟踪器
    with common.RewardTracker(writer, stop_reward=10000.0, logger=logger) as reward_tracker:
        while True:
            frame_idx += 1
            # 从经验池中获取一次样本存放在缓存中
            buffer.populate(1)
            # 更新epsilon值
            epsilon_tracker.frame(frame_idx)

            # 这个操作是将所有采集的激励pop出缓存并清空缓存
            # 这里清空仅仅只清空exp_source中的记录，在buffer中的记录并没有清空
            new_rewards = exp_source.pop_total_rewards()
            if new_rewards:
                # 这里主要是判断最近的一次激励是否达到了目标值，达到则停止训练
                if reward_tracker.reward(new_rewards[0], frame_idx, selector.epsilon):
                    # 将当前采集的首个激励，帧总数传入到reward_tracker
                    # 如果平均激励值已经达到了所要的目标值，则break跳出循环
                    break

            # 检查经验池中的样本长度是否达到了目标的长度大小，
            if len(buffer) < 20000:
                continue

            # 清空优化器的梯度
            optimizer.zero_grad()
            # 从经验缓冲区中采集batch_size大小的样本
            batch = buffer.sample(128)
            # 计算损失值并更新神经网路
            loss_v = common.calc_loss_dqn(batch, net, tgt_net.target_model, gamma=0.99, device=device)
            loss_v.backward()
            optimizer.step()
            scheduler.step()

            if frame_idx % 1000 == 0:
                tgt_net.sync()
            if frame_idx % 5000 == 0:
                # Test the model
                test_reward = test_model(test_env, net, device=device, episodes=5)
                print(f"Test reward: {test_reward:.2f}")
                common.save_best_model(test_reward, net.state_dict(), save_path, "dqn-basic-best", keep_best=10)


                checkpoint = {
                    'net': net.state_dict(),
                    'tgt_net': tgt_net.target_model.state_dict(),
                    'optimizer': optimizer.state_dict(),
                    'frame_idx': frame_idx,
                    'scheduler': scheduler.state_dict()
                }
                common.save_checkpoints(frame_idx, checkpoint, save_path, "dqn-basic", keep_last=5)
